#pragma once

#include <cstddef>
#include <mutex>
#include <vector>
#include <memory>
#include <array>
#include <list>
#include <unordered_map>
#include <cassert>

namespace SeanNetwork {

/**
 * @brief 固定大小块内存池
 * 
 * 提供高效的同类型固定大小内存块分配和回收
 */
template<size_t BlockSize>
class FixedSizeMemoryPool {
public:
    FixedSizeMemoryPool(size_t initialBlocks = 32)
        : m_FreeBlocks(nullptr), m_AllocatedChunks(0) {
        // 分配初始内存块
        AllocateChunk(initialBlocks);
    }
    
    ~FixedSizeMemoryPool() {
        // 释放所有分配的内存
        for (auto& chunk : m_Chunks) {
            ::operator delete(chunk);
        }
    }
    
    // 禁止复制
    FixedSizeMemoryPool(const FixedSizeMemoryPool&) = delete;
    FixedSizeMemoryPool& operator=(const FixedSizeMemoryPool&) = delete;
    
    /**
     * @brief 分配一个内存块
     * @return 指向新分配内存块的指针
     */
    void* Allocate() {
        std::lock_guard<std::mutex> lock(m_Mutex);
        
        if (m_FreeBlocks == nullptr) {
            // 没有空闲块，分配新的块组
            const size_t expansionSize = m_AllocatedChunks > 0 ? m_AllocatedChunks : 32;
            AllocateChunk(expansionSize);
        }
        
        // 获取一个空闲块
        void* block = m_FreeBlocks;
        m_FreeBlocks = *reinterpret_cast<void**>(block);
        
        return block;
    }
    
    /**
     * @brief 释放内存块
     * @param block 要释放的内存块指针
     */
    void Deallocate(void* block) {
        if (!block) return;
        
        std::lock_guard<std::mutex> lock(m_Mutex);
        
        // 将块放回空闲列表
        *reinterpret_cast<void**>(block) = m_FreeBlocks;
        m_FreeBlocks = block;
    }
    
    /**
     * @brief 获取当前已分配的块数
     */
    size_t GetAllocatedChunks() const {
        std::lock_guard<std::mutex> lock(const_cast<std::mutex&>(m_Mutex));
        return m_AllocatedChunks;
    }
    
private:
    /**
     * @brief 分配一组新的内存块
     * @param count 要分配的块数
     */
    void AllocateChunk(size_t count) {
        // 分配一块足够大的内存，可容纳count个块
        size_t chunkSize = BlockSize * count;
        auto* chunk = static_cast<char*>(::operator new(chunkSize));
        m_Chunks.push_back(chunk);
        
        // 将新块链接到空闲列表
        void* firstBlock = chunk;
        void* lastBlock = static_cast<char*>(chunk) + BlockSize * (count - 1);
        
        for (size_t i = 0; i < count - 1; ++i) {
            void* currentBlock = static_cast<char*>(chunk) + BlockSize * i;
            void* nextBlock = static_cast<char*>(chunk) + BlockSize * (i + 1);
            *reinterpret_cast<void**>(currentBlock) = nextBlock;
        }
        
        // 最后一个块的下一个指针为当前空闲列表
        *reinterpret_cast<void**>(lastBlock) = m_FreeBlocks;
        m_FreeBlocks = firstBlock;
        
        m_AllocatedChunks += count;
    }
    
    mutable std::mutex m_Mutex;       // 保护资源互斥锁
    void* m_FreeBlocks;               // 空闲块链表
    std::vector<void*> m_Chunks;      // 已分配的内存块组
    size_t m_AllocatedChunks;         // 已分配的块总数
};

/**
 * @brief 分层内存池
 * 
 * 根据请求的大小使用不同的固定大小内存池
 */
class TieredMemoryPool {
public:
    TieredMemoryPool()
        : m_SmallBlocksPool(128),   // 128字节块
          m_MediumBlocksPool(1024), // 1KB块
          m_LargeBlocksPool(4096)   // 4KB块
    {}
    
    ~TieredMemoryPool() = default;
    
    // 禁止复制
    TieredMemoryPool(const TieredMemoryPool&) = delete;
    TieredMemoryPool& operator=(const TieredMemoryPool&) = delete;
    
    /**
     * @brief 分配指定大小的内存
     * @param size 请求的内存大小
     * @return 指向已分配内存的指针，如果分配失败返回nullptr
     */
    void* Allocate(size_t size) {
        if (size == 0) return nullptr;
        
        std::lock_guard<std::mutex> lock(m_Mutex);
        
        void* memory = nullptr;
        
        if (size <= 128) {
            memory = m_SmallBlocksPool.Allocate();
        } else if (size <= 1024) {
            memory = m_MediumBlocksPool.Allocate();
        } else if (size <= 4096) {
            memory = m_LargeBlocksPool.Allocate();
        } else {
            // 对于超大块，直接使用标准分配器
            memory = ::operator new(size);
            m_OversizedBlocks[memory] = size;
        }
        
        // 记录分配的内存大小类别
        if (memory) {
            m_AllocationSizes[memory] = size;
        }
        
        return memory;
    }
    
    /**
     * @brief 释放先前分配的内存
     * @param memory 要释放的内存指针
     */
    void Deallocate(void* memory) {
        if (!memory) return;
        
        std::lock_guard<std::mutex> lock(m_Mutex);
        
        // 查找分配的大小
        auto it = m_AllocationSizes.find(memory);
        if (it == m_AllocationSizes.end()) {
            // 未知的内存指针，忽略
            return;
        }
        
        size_t size = it->second;
        
        if (size <= 128) {
            m_SmallBlocksPool.Deallocate(memory);
        } else if (size <= 1024) {
            m_MediumBlocksPool.Deallocate(memory);
        } else if (size <= 4096) {
            m_LargeBlocksPool.Deallocate(memory);
        } else {
            // 对于超大块，直接使用标准释放
            ::operator delete(memory);
            m_OversizedBlocks.erase(memory);
        }
        
        // 从记录中移除
        m_AllocationSizes.erase(it);
    }
    
    /**
     * @brief 获取当前内存池的统计信息
     * @return 包含各种大小分配统计的字符串
     */
    std::string GetStats() const {
        std::string stats;
        
        std::lock_guard<std::mutex> lock(const_cast<std::mutex&>(m_Mutex));
        
        size_t smallCount = m_SmallBlocksPool.GetAllocatedChunks();
        size_t mediumCount = m_MediumBlocksPool.GetAllocatedChunks();
        size_t largeCount = m_LargeBlocksPool.GetAllocatedChunks();
        size_t oversizedCount = m_OversizedBlocks.size();
        
        stats += "Memory Pool Stats:\n";
        stats += "  Small Blocks (<=128B): " + std::to_string(smallCount) + " chunks\n";
        stats += "  Medium Blocks (<=1KB): " + std::to_string(mediumCount) + " chunks\n";
        stats += "  Large Blocks (<=4KB): " + std::to_string(largeCount) + " chunks\n";
        stats += "  Oversized Blocks (>4KB): " + std::to_string(oversizedCount) + " blocks\n";
        stats += "  Total Allocations: " + std::to_string(m_AllocationSizes.size()) + " blocks\n";
        
        return stats;
    }
    
private:
    mutable std::mutex m_Mutex; // 保护公共资源的互斥锁
    
    // 不同大小的内存池
    FixedSizeMemoryPool<128> m_SmallBlocksPool;   // 小块 (<=128B)
    FixedSizeMemoryPool<1024> m_MediumBlocksPool; // 中块 (<=1KB)
    FixedSizeMemoryPool<4096> m_LargeBlocksPool;  // 大块 (<=4KB)
    
    // 记录每个分配的内存块的大小
    std::unordered_map<void*, size_t> m_AllocationSizes;
    
    // 记录特别大的内存块
    std::unordered_map<void*, size_t> m_OversizedBlocks;
};

/**
 * @brief 分配器适配器，用于将内存池用作STL容器的分配器
 */
template<typename T>
class MemoryPoolAllocator {
public:
    using value_type = T;
    using pointer = T*;
    using const_pointer = const T*;
    using reference = T&;
    using const_reference = const T&;
    using size_type = std::size_t;
    using difference_type = std::ptrdiff_t;
    
    template<typename U>
    struct rebind {
        using other = MemoryPoolAllocator<U>;
    };
    
    MemoryPoolAllocator() noexcept
        : m_MemoryPool(GetGlobalMemoryPool()) {}
    
    template<typename U>
    MemoryPoolAllocator(const MemoryPoolAllocator<U>&) noexcept
        : m_MemoryPool(GetGlobalMemoryPool()) {}
    
    // 分配内存
    T* allocate(size_t n) {
        if (n == 0) return nullptr;
        
        void* memory = m_MemoryPool->Allocate(n * sizeof(T));
        if (!memory) throw std::bad_alloc();
        
        return static_cast<T*>(memory);
    }
    
    // 释放内存
    void deallocate(T* p, size_t) noexcept {
        if (p) m_MemoryPool->Deallocate(p);
    }
    
    // 获取全局内存池实例
    static TieredMemoryPool* GetGlobalMemoryPool() {
        static TieredMemoryPool globalPool;
        return &globalPool;
    }
    
private:
    TieredMemoryPool* m_MemoryPool;
};

// 检查两个分配器是否相等
template<typename T1, typename T2>
bool operator==(const MemoryPoolAllocator<T1>&, const MemoryPoolAllocator<T2>&) noexcept {
    return true;
}

template<typename T1, typename T2>
bool operator!=(const MemoryPoolAllocator<T1>&, const MemoryPoolAllocator<T2>&) noexcept {
    return false;
}

} // namespace SeanNetwork 